from typing import Callable, Any


class NoServiceError(Exception):
    """无服务错误"""
    pass


class Route(object):
    """
    rpc服务路由
    """

    def __init__(self):
        self.key2Func = {}

    def add(self, key: str, func: Callable):
        """注册服务"""
        if key in self.key2Func:
            raise KeyError(key)
        self.key2Func[key] = func

    def route(self, key: str = "") -> Callable:
        """装饰器方式注册服务"""

        def decorator(func):
            _key = key if key else func.__name__
            if _key in self.key2Func:
                raise KeyError(_key)

            self.key2Func[_key] = func

            return func

        return decorator

    async def call(self, key: str, *args, **kwargs) -> Any:
        """调用服务"""
        func = self.key2Func.get(key)
        if not func:
            raise NoServiceError(key)

        return await func(*args, **kwargs)
